"""
PyG版本的图池化操作，与dgl.nn.pytorch.AvgPooling和dgl.nn.pytorch.MaxPooling类兼容
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import torch.nn as nn
from torch_geometric.nn import global_mean_pool, global_max_pool
from ops.src.graph.graph import Graph


class AvgPooling(nn.Module):
    """
    平均池化操作，与dgl.nn.pytorch.AvgPooling类兼容
    """
    
    def __init__(self):
        super().__init__()
    
    def forward(self, graph: Graph, feat: torch.Tensor):
        """
        对图中的节点特征进行平均池化
        
        Args:
            graph: 输入图
            feat: 节点特征，形状为[num_nodes, feat_dim, ...]
            
        Returns:
            池化后的特征，形状为[batch_size, feat_dim, ...]
        """
        print("===================AvgPooling===================")
        # 获取批处理信息
        batch = graph.batch if hasattr(graph, 'batch') else None
        
        # 调试信息已关闭
        
        if batch is None:
            # 这种情况现在不应该发生，因为我们修复了batch函数
            raise RuntimeError("❌ 批处理信息缺失！图的batch属性未正确设置")
        
        # 🔧 确保设备一致性
        if batch.device != feat.device:
            batch = batch.to(feat.device)
        
        # 获取特征维度
        feat_shape = feat.shape
        if len(feat_shape) > 2:
            # 处理多维特征
            feat_dim = feat_shape[1:]
            feat_flat = feat.view(feat_shape[0], -1)
            # 使用PyG的global_mean_pool进行池化
            pooled = global_mean_pool(feat_flat, batch)
            # 恢复原始维度
            pooled = pooled.view(-1, *feat_dim)
        else:
            # 处理二维特征
            pooled = global_mean_pool(feat, batch)
        
        return pooled


class MaxPooling(nn.Module):
    """
    最大池化操作，与dgl.nn.pytorch.MaxPooling类兼容
    """
    print("===================MaxPooling===================")
    def __init__(self):
        super().__init__()
    
    def forward(self, graph: Graph, feat: torch.Tensor):
        """
        对图中的节点特征进行最大池化
        
        Args:
            graph: 输入图
            feat: 节点特征，形状为[num_nodes, feat_dim, ...]
            
        Returns:
            池化后的特征，形状为[batch_size, feat_dim, ...]
        """
        # 获取批处理信息
        batch = graph.batch if hasattr(graph, 'batch') else None
        if batch is None:
            raise RuntimeError("❌ 批处理信息缺失！图的batch属性未正确设置")
        
        # 🔧 确保设备一致性
        if batch.device != feat.device:
            batch = batch.to(feat.device)
        
        # 获取特征维度
        feat_shape = feat.shape
        if len(feat_shape) > 2:
            # 处理多维特征
            feat_dim = feat_shape[1:]
            feat_flat = feat.view(feat_shape[0], -1)
            # 使用PyG的global_max_pool进行池化
            pooled = global_max_pool(feat_flat, batch)
            # 恢复原始维度
            pooled = pooled.view(-1, *feat_dim)
        else:
            # 处理二维特征
            pooled = global_max_pool(feat, batch)
        
        return pooled 